Skip to content

Add MemoryVisualizationCallback #734

Merged
tomwhite merged 2 commits intomainfrom
memory-usage-timeline
Jun 21, 2025
Merged

Add MemoryVisualizationCallback #734
tomwhite merged 2 commits intomainfrom
memory-usage-timeline

Conversation

@tomwhite
Copy link
Copy Markdown
Member

This is based on @TomNicholas' work in https://gist.github.com/TomNicholas/c6a28f7c22c6981f75bce280d3e28283, the main difference being that the x-axis is task (number) rather than time, which gives a less distorted plot. (Using time tends to overplot many tasks on top of each other as many tend to start or finish at about the same time.)

Here's an example for the ERA5 rechunking I did in #650:

era5-mem-usage

Interestingly, the last stage of rechunking seems to have exceeded the projected memory - but not the allowed memory, so it didn't fail. It would be worth looking into this to see if the projected memory calculation can be improved though.

@TomNicholas
Copy link
Copy Markdown
Member

This looks great!

nit: Is "max peak measured" a tautology?

@tomwhite
Copy link
Copy Markdown
Member Author

This looks great!

Thanks!

nit: Is "max peak measured" a tautology?

I don't think so. Each task has a peak memory usage, and we find the maximum across all tasks.

But maybe it's fine to just say "max measured" or "max usage".

@tomwhite tomwhite merged commit 5d4e76d into main Jun 21, 2025
19 of 20 checks passed
@tomwhite tomwhite deleted the memory-usage-timeline branch June 21, 2025 07:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants